package com.github.cyf.rateLimit.annotation;

import java.util.concurrent.TimeUnit;

/**
 * @author chenyifan
 * @create 2025-02-05 10:47
 */
public class LeakyBucketLimiter {
    /**
     * 桶的容量
     */
    private final int capacity;
    /**
     * 漏桶的漏出速率，单位时间内漏出水的数量
     */
    private final int leakRate;
    /**
     * 当前桶中的水量
     */
    private int water = 0;
    /**
     * 上次漏水的时间
     */
    private long lastLeakTime = System.currentTimeMillis();

    private LeakyBucketLimiter(int capacity, int leakRate) {
        this.capacity = capacity;
        this.leakRate = leakRate;
    }

    public static LeakyBucketLimiter create(int capacity, int leakRate) {
        return new LeakyBucketLimiter(capacity, leakRate);
    }

    public boolean tryAcquire() {
        long currentTime = System.currentTimeMillis();
        synchronized (this) {
            // 计算上次漏水到当前时间的时间间隔
            long timeGap = currentTime - lastLeakTime;
            // 如果时间间隔大于1秒，重新计算桶中的水量
            if (timeGap >= TimeUnit.SECONDS.toMillis(1)) {
                // 计算漏出的水量
                int leakAmount = (int) (timeGap / TimeUnit.SECONDS.toMillis(1) * leakRate);
                // 更新桶中的水量，但不能低于0
                water = Math.max(0, water - leakAmount);
                lastLeakTime = currentTime;
            }
            // 判断桶中的水量是否小于容量，如果是则可以继续添加水（相当于获取到令牌）
            if (water < capacity) {
                water++;
                return true;
            }
        }
        return false;
    }
}